import os

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, TensorDataset
from torch.nn.utils.rnn import pad_sequence
# 创建滑动窗口的数据集
# acce: ax, ay, az, gx, gy, gz (N, 6) angle: (N,) 对齐时间戳
def data_align(imu_data, gps_data):
    imu_time = imu_data[:,0]
    gps_time = gps_data[:,0]
    imu_slices = []

    for i in range(len(gps_time)-1):
        t1 = gps_time[i]
        t2 = gps_time[i+1]

        idex_t1 = np.searchsorted(imu_time, t1, side='left')
        idex_t2 = np.searchsorted(imu_time, t2, side='right')

        idx_min = max(0, idex_t1 - 50)
        idx_max = min(len(imu_time), idex_t2 + 50)

        temp_slice = imu_data[idx_min:idx_max]
        imu_slices.append(temp_slice)
    return imu_slices, gps_data

def read_data_file(head):
    data = os.listdir(head)
    imu_list = []
    vi_list = []
    for d in data:
        if 'imu1' in d:
            imu_t = pd.read_csv(os.path.join(head, d)).to_numpy()
            imu_list.append(imu_t)
        if 'vi1' in d:
            vi_t = pd.read_csv(os.path.join(head, d)).to_numpy()
            vi_list.append(vi_t)
            break
    # 时间戳一致，直接读一个就行
    return imu_list, vi_list

def get_dataloader(batch_size):

    imu_list = []
    vi_list = []
    dir = '../dataset/oxford_handheld'
    for i in range(1,2):
        head = dir + '/' + 'data{}'.format(str(i)) + '/' + 'syn'
        imu_list_t, vi_list_t = read_data_file(head)
        imu_list += imu_list_t
        vi_list += vi_list_t
    imu_np = None
    vi_np = None
    for i, v in zip(imu_list, vi_list):
        if imu_np is None and vi_np is None:
            imu_np = i
            vi_np = v
        else:
            imu_np = np.concatenate((imu_np, i), axis=0)
            vi_np = np.concatenate((vi_np, v), axis=0)
    # rx ry rz ax ay az
    imu_np = imu_np[:,[1,2,3,10,11,12]]
    vi_np = vi_np[:,[2,3]]

    relative_polar_coords = []
    relative_polar_coords.append((0, 0))
    for i in range(1, len(vi_np)):

        dx = vi_np[i][0] - vi_np[i-1][0]
        dy = vi_np[i][1] - vi_np[i-1][1]

        r = np.sqrt(dx**2 + dy**2)
        angle = np.arctan2(dy, dx)
        relative_polar_coords.append((r, angle))
    relative_polar_coords = np.array(relative_polar_coords)
    divi = int(0.1 * len(relative_polar_coords))
    imu_tensor = torch.tensor(imu_np, dtype=torch.float32)
    imu_tensor = (imu_tensor - torch.min(imu_tensor)) / (torch.max(imu_tensor) - torch.min(imu_tensor))
    relative_polar_coords_tensor = torch.tensor(relative_polar_coords, dtype=torch.float32)
    test_set = TensorDataset(imu_tensor[:divi,:], relative_polar_coords_tensor[:divi,1])
    train_set = TensorDataset(imu_tensor[divi:,:], relative_polar_coords_tensor[divi:,1])
    print(len(test_set))
    print(len(train_set))
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=False)
    return train_loader, test_loader
#train, test = get_dataloader(1000)
